from torch import nn
from mmcv.cnn.resnet import constant_init, kaiming_init
import torch
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import GCnet原始版本,GCnet改进版本,NLNet消融实验,Simplified_NLNet消融实验
import matplotlib.pyplot as plt

if __name__ == '__main__':

    # 设置随机种子以保证结果可重现
    torch.manual_seed(123)
    torch.cuda.manual_seed(123)

    # 定义数据预处理
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])

    batch_size=128

    # 加载训练数据集和测试数据集
    trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
    testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)

    # 定义模型,在此处选择模型

    model = GCnet原始版本.GC_all_resnet50()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)

    # 设置训练参数
    num_epochs = 2
    log_interval = 391

    # 创建记录Loss和Accuracy的列表
    train_loss_list = []
    train_acc_list = []
    test_loss_list = []
    test_acc_list = []

    # 创建TensorBoard可视化对象
    writer = SummaryWriter()

    # 训练模型
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0

        for batch_idx, (inputs, targets) in enumerate(tqdm(trainloader)):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            # 统计训练信息
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

            # 打印训练信息
            if (batch_idx + 1) % log_interval == 0:
                acc = 100. * correct / total
                avg_loss = running_loss / log_interval
                print(f"Epoch [{epoch + 1}/{num_epochs}], Step [{batch_idx + 1}/{len(trainloader)}], Loss: {avg_loss:.4f}, Accuracy: {acc:.2f}%")
                train_loss_list.append(avg_loss)
                train_acc_list.append(acc)
                writer.add_scalar("Loss/train", avg_loss, epoch * len(trainloader) + batch_idx)
                writer.add_scalar("Accuracy/train", acc, epoch * len(trainloader) + batch_idx)

                running_loss = 0.0
                correct = 0
                total = 0

        # 在测试集上评估模型
        model.eval()
        test_loss = 0.0
        test_correct = 0
        test_total = 0

        with torch.no_grad():
            for inputs, targets in testloader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                test_total += targets.size(0)
                test_correct += predicted.eq(targets).sum().item()

        test_acc = 100. * test_correct / test_total
        avg_test_loss = test_loss / len(testloader)
        print(f"Epoch [{epoch + 1}/{num_epochs}], Test Loss: {avg_test_loss:.4f}, Test Accuracy: {test_acc:.2f}%")
        writer.add_scalar("Loss/test", avg_test_loss, epoch)
        writer.add_scalar("Accuracy/test", test_acc, epoch)
        test_loss_list.append(avg_test_loss)
        test_acc_list.append(test_acc)

    # 保存模型
    torch.save(model.state_dict(), "gcnet_cifar100.pth")
    writer.close()

    epochs = range(1, num_epochs + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_loss_list, 'b-', label='Train')
    plt.plot(epochs, test_loss_list, 'r-', label='Test')
    plt.title('Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs, train_acc_list, 'b-', label='Train')
    plt.plot(epochs, test_acc_list, 'r-', label='Test')
    plt.title('Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()